from typing import List, Tuple
from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import PGVector
from langchain.document_loaders import TextLoader
from langchain.docstore.document import Document
from langchain_community.embeddings import HuggingFaceEmbeddings
import warnings
warnings.filterwarnings("ignore")

loader = TextLoader("../demo.txt", encoding="utf-8")
documents = loader.load()
text_splitter = CharacterTextSplitter(
    separator="\n\n",
    chunk_size=75,
    chunk_overlap=30
)
docs = text_splitter.split_documents(documents)

embeddings = HuggingFaceEmbeddings(model_name='BAAI/bge-base-zh')

import os

CONNECTION_STRING = PGVector.connection_string_from_db_params(
    driver=os.environ.get("PGVECTOR_DRIVER", "psycopg2"),
    host=os.environ.get("PGVECTOR_HOST", "localhost"),
    port=int(os.environ.get("PGVECTOR_PORT", "5432")),
    database=os.environ.get("PGVECTOR_DATABASE", "postgres"),
    user=os.environ.get("PGVECTOR_USER", "postgres"),
    password=os.environ.get("PGVECTOR_PASSWORD", "123"),
)

db = PGVector.from_documents(
    embedding=embeddings,
    documents=docs,
    collection_name="state_of_the_union",
    connection_string=CONNECTION_STRING,
)

query = "微笑明天慈善基金志愿者算社会实践活动省级奖励吗？"
docs_with_score: List[Tuple[Document, float]] = db.similarity_search_with_score(query)


for doc, score in docs_with_score:
    print("-" * 80)
    print("Score: ", score)
    print(doc.page_content)
    print("-" * 80)


llm = Ollama(model="qwen2.5")
prompt = ChatPromptTemplate.from_template("""仅根据提供的上下文回答以下问题:

<context>
{context}
</context>

Question: {input}""")

document_chain = prompt | llm


print(document_chain.invoke({
    "input": "微笑明天慈善基金志愿者算社会实践活动省级奖励吗？",
    "context": docs_with_score[0][0].page_content
}))
